import os
import glob
from datetime import datetime
from tqdm import tqdm
from pathlib import Path
import json
import base64
from PIL import Image
import io

from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter

# 设置 API 密钥
os.environ["NVIDIA_API_KEY"] = "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"

# 初始化嵌入模型
embedder = NVIDIAEmbeddings(model="NV-Embed-QA")

def create_product_embeddings(product_dir="./zh_data/products", embedding_dir="./zh_data/embeddings", chunk_size=400):
    """
    根据商品目录下的图片和描述文件创建词嵌入
    
    参数:
        product_dir: 商品目录
        embedding_dir: 嵌入保存目录
        chunk_size: 文本分块大小
        chunk_overlap: 文本块重叠大小
    
    返回:
        嵌入文件路径
    """
    # 确保目录存在
    os.makedirs(embedding_dir, exist_ok=True)
    os.makedirs(product_dir, exist_ok=True)
    
    # 检查商品目录是否存在
    if not os.path.exists(product_dir):
        print(f"错误: 目录 {product_dir} 不存在")
        return None
    
    # 初始化嵌入模型
    embedder = NVIDIAEmbeddings(model="NV-Embed-QA")
    
    # 读取商品数据
    data = []
    sources = []
    
    # 获取所有商品JSON文件
    product_files = glob.glob(os.path.join(product_dir, "**/*.json"), recursive=True)
    
    for product_file in product_files:
        try:
            with open(product_file, encoding="utf-8") as f:
                product_data = json.load(f)
            
            # 处理列表格式的商品数据
            if isinstance(product_data, list):
                for item in product_data:
                    # 构建商品描述文本
                    product_text = f"商品ID: {item.get('id', '')}\n商品名称: {item.get('name', '')}\n商品类别: {item.get('category', '')}\n商品价格: {item.get('price', '')}\n商品描述: {item.get('description', '')}"
                    
                    data.append(product_text)
                    sources.append(product_file)
            else:
                # 处理单个商品数据
                product_text = f"商品ID: {product_data.get('id', '')}\n商品名称: {product_data.get('name', '')}\n商品类别: {product_data.get('category', '')}\n商品价格: {product_data.get('price', '')}\n商品描述: {product_data.get('description', '')}"
                
                data.append(product_text)
                sources.append(product_file)

            print(product_text)
        except Exception as e:
            print(f"处理商品文件 {product_file} 时出错: {str(e)}")
    
    if not data:
        print("没有找到有效的商品数据")
        return None
    
    print(f"读取了 {len(data)} 个商品")
    
    # 分割文本
    text_splitter = CharacterTextSplitter(chunk_size=chunk_size, separator=" ")
    docs = []
    metadatas = []
    
    for i, d in enumerate(data):
        splits = text_splitter.split_text(d)
        docs.extend(splits)
        metadatas.extend([{"source": sources[i]}] * len(splits))
    
    print(f"分割后共有 {len(docs)} 个文本块")
    
    # 创建向量存储
    try:
        store = FAISS.from_texts(docs, embedder, metadatas=metadatas)

        # 保存向量存储
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        embedding_path = os.path.join(embedding_dir, f"product_embedding_{timestamp}")
        store.save_local(embedding_path)
        
        print(f"商品向量存储已保存到 {embedding_path}")
        return embedding_path
    except Exception as e:
        print(f"创建向量存储时出错: {str(e)}")
        return None

# 创建示例商品数据
def create_sample_products(product_dir="./zh_data/products"):
    """
    创建示例商品数据
    
    参数:
        product_dir: 商品目录
    """
    # 确保目录存在
    os.makedirs(product_dir, exist_ok=True)
    
    # 创建示例商品类别
    categories = ["手机", "电脑", "服装", "食品"]
    
    for category in categories:
        category_path = os.path.join(product_dir, category)
        os.makedirs(category_path, exist_ok=True)
        
        # 创建示例商品
        if category == "手机":
            products = [
                {
                    "id": "phone001",
                    "name": "iPhone 15 Pro",
                    "category": "手机",
                    "price": "8999元",
                    "description": "iPhone 15 Pro 采用 A17 Pro 芯片，配备 6.1 英寸 Super Retina XDR 显示屏，支持 ProMotion 技术，最高可达 120Hz 刷新率。相机系统包括 4800 万像素主摄、1200 万像素超广角镜头和 1200 万像素长焦镜头。支持 5G 网络，电池容量为 3200mAh。",
                    "image_path": "./zh_data/images/iphone15pro.jpg"
                },
                {
                    "id": "phone002",
                    "name": "华为 Mate 60 Pro",
                    "category": "手机",
                    "price": "6999元",
                    "description": "华为 Mate 60 Pro 采用麒麟 9000S 芯片，配备 6.82 英寸 LTPO OLED 显示屏，支持 120Hz 刷新率。相机系统包括 5000 万像素主摄、4800 万像素超广角镜头和 4800 万像素长焦镜头。支持 5G 网络，电池容量为 5000mAh。",
                    "image_path": "./zh_data/images/mate60pro.jpg"
                }
            ]
        elif category == "电脑":
            products = [
                {
                    "id": "laptop001",
                    "name": "MacBook Pro 14",
                    "category": "电脑",
                    "price": "14999元",
                    "description": "MacBook Pro 14 采用 M2 Pro 或 M2 Max 芯片，配备 14.2 英寸 Liquid Retina XDR 显示屏，支持 ProMotion 技术，最高可达 120Hz 刷新率。内存最高可达 96GB，存储容量最高可达 8TB。电池容量为 70Wh，支持快速充电。",
                    "image_path": "./zh_data/images/macbookpro14.jpg"
                },
                {
                    "id": "laptop002",
                    "name": "联想 ThinkPad X1 Carbon",
                    "category": "电脑",
                    "price": "12999元",
                    "description": "联想 ThinkPad X1 Carbon 采用 Intel Core i7-1260P 处理器，配备 14 英寸 WQXGA 显示屏，支持 90Hz 刷新率。内存最高可达 32GB，存储容量最高可达 2TB。电池容量为 57Wh，支持快速充电。",
                    "image_path": "./zh_data/images/thinkpadx1.jpg"
                }
            ]
        elif category == "服装":
            products = [
                {
                    "id": "clothes001",
                    "name": "优衣库男士休闲衬衫",
                    "category": "服装",
                    "price": "199元",
                    "description": "优衣库男士休闲衬衫采用优质棉质面料，舒适透气，适合日常穿着。有多种颜色可选，包括白色、蓝色、灰色等。尺码从 S 到 XXL 都有。",
                    "image_path": "./zh_data/images/uniqlo_shirt.jpg"
                },
                {
                    "id": "clothes002",
                    "name": "ZARA 女士连衣裙",
                    "category": "服装",
                    "price": "399元",
                    "description": "ZARA 女士连衣裙采用轻盈面料，设计简约时尚，适合各种场合穿着。有多种颜色和图案可选，尺码从 XS 到 L 都有。",
                    "image_path": "./zh_data/images/zara_dress.jpg"
                }
            ]
        elif category == "食品":
            products = [
                {
                    "id": "food001",
                    "name": "三只松鼠坚果大礼包",
                    "category": "食品",
                    "price": "99元",
                    "description": "三只松鼠坚果大礼包包含多种坚果，包括夏威夷果、腰果、开心果、碧根果等。每袋独立包装，保持新鲜。适合送礼或自己食用。",
                    "image_path": "./zh_data/images/three_squirrels.jpg"
                },
                {
                    "id": "food002",
                    "name": "百事可乐",
                    "category": "食品",
                    "price": "3元",
                    "description": "百事可乐是一种碳酸饮料，口感清爽，气泡丰富。有多种规格可选，包括 330ml、500ml、1L、2L 等。适合各种场合饮用。",
                    "image_path": "./zh_data/images/pepsi.jpg"
                }
            ]
        
        # 保存商品数据
        for product in products:
            product_file = os.path.join(category_path, f"{product['id']}.json")
            with open(product_file, "w", encoding="utf-8") as f:
                json.dump(product, f, ensure_ascii=False, indent=4)
            print(f"已创建商品文件: {product_file}")

# 创建示例商品图片目录
def create_sample_images(image_dir="./zh_data/images"):
    """
    创建示例商品图片目录
    
    参数:
        image_dir: 图片目录
    """
    # 确保目录存在
    os.makedirs(image_dir, exist_ok=True)
    
    # 创建示例图片
    # 这里只是创建目录，实际图片需要用户自己添加
    print(f"请将商品图片添加到 {image_dir} 目录中")
    print("需要的图片文件:")
    print("- iphone15pro.jpg")
    print("- mate60pro.jpg")
    print("- macbookpro14.jpg")
    print("- thinkpadx1.jpg")
    print("- uniqlo_shirt.jpg")
    print("- zara_dress.jpg")
    print("- three_squirrels.jpg")
    print("- pepsi.jpg")

if __name__ == "__main__":
    # 创建示例商品图片目录
    # create_sample_images()
    
    # 创建示例商品数据
    # create_sample_products()
    
    # 创建商品嵌入
    embedding_path = create_product_embeddings()
    
    if embedding_path:
        print(f"商品嵌入已创建: {embedding_path}")
    else:
        print("创建商品嵌入失败") 